import angr
import claripy
import networkx as nx
import logging
import sys

import angr.knowledge_base

sys.path.append('src')

from tasks.lsbvd.feature.const_value import buffer_overflow_funcs, buffer_overflow_sink_funcs, buffer_overflow_source_funcs
from angr.analyses.ddg import DDG

class MIPSTaintTracer:
    def __init__(self, proj, cfg):
        self.proj = proj
        self.cfg = cfg

    def trace_parameter_origin(self, function_path):
        logging.getLogger('angr').setLevel(logging.INFO)
        if not function_path or len(function_path) < 2:
            return False, None

        first_func_addr = function_path[0]
        second_last_func_addr = function_path[-2]

        init_state = self.proj.factory.blank_state(
            addr=first_func_addr,
            add_options={
                angr.options.ZERO_FILL_UNCONSTRAINED_REGISTERS,
                angr.options.ZERO_FILL_UNCONSTRAINED_MEMORY
            }
        )
        params = {
            'a0': claripy.BVS('arg0', 32),
            'a1': claripy.BVS('arg1', 32),
            'a2': claripy.BVS('arg2', 32),
            'a3': claripy.BVS('arg3', 32)
        }
        init_state.regs.a0 = params['a0']
        init_state.regs.a1 = params['a1']
        init_state.regs.a2 = params['a2']
        init_state.regs.a3 = params['a3']

        simgr = self.proj.factory.simulation_manager(init_state)
        simgr.use_technique(angr.exploration_techniques.LengthLimiter(max_length=10000))

        for i, func_addr in enumerate(function_path):
            func = self.cfg.kb.functions.get(func_addr)
            if func is None:
                return False, func_addr
            exit_points = [b.addr for b in func.endpoints] + [b.addr for b in func.ret_sites]
            simgr.explore(
                find=exit_points,
                num_find=1,
            )
            found_states = simgr.stashes.get('found', [])
            if found_states:
                current_state = found_states[0].copy()
                simgr = self.proj.factory.simulation_manager(current_state)
                if i < len(function_path) - 1:
                    next_func_addr = function_path[i+1]
                    simgr.active[0].ip = next_func_addr
            else:
                return False, func_addr
            if func_addr == second_last_func_addr:
                return True, function_path

        return False, function_path[-1]

if __name__ == "__main__":
    import time
    import json
    import os
    import pickle
    from tasks.lsbvd.feature.static_cfg import BigBinaryCfgBuilder

    region_split_path = 'caches_data/amd.txtbin_partitions.json'
    region_json = json.load(open(region_split_path, 'r'))['partitions']

    input_path = "caches_data/nokia_path/output.txt"
    output_path = "caches_data/nokia_path/success_paths.txt"
    fail_path = "caches_data/nokia_path/failed_funcs.txt"

    bbinary_builder = BigBinaryCfgBuilder(
        binary_path='caches_data/ine.txtbin',
        caches_dir='caches_data/nokia_cfg',
        symbols_json_path='caches_data/symbols_export.json'
    )

    cfgs = bbinary_builder.load_partial_cfg(regions=region_json)
    cfg = None
    for e in cfgs:
        cfg = e
        break

    proj = bbinary_builder.proj
    tracer = MIPSTaintTracer(proj, cfg)

    with open(input_path, "r") as fin, \
         open(output_path, "w") as fout, \
         open(fail_path, "w") as ffail:
        for line in fin:
            line = line.strip()
            if not line:
                continue

            try:
                function_path = [int(addr.strip()) for addr in line.split(",") if addr.strip()]
            except Exception:
                continue
            print(f"\nTrying FUNCTION_PATH: {[hex(addr) for addr in function_path]}")
            result, info = tracer.trace_parameter_origin(function_path)
            if result:
                fout.write(",".join([hex(addr) for addr in function_path]) + "\n")
                print("溯源成功，已写入结果文件。")
            else:
                ffail.write(hex(info) + "\n")
                print(f"溯源失败，已记录失败函数地址：{hex(info)}")
    print("全部处理完成。")